{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a4d5d1bc-df65-47ae-8e2a-d1fe6218fb3e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-31T03:04:54.663930Z",
     "iopub.status.busy": "2022-05-31T03:04:54.663218Z",
     "iopub.status.idle": "2022-05-31T03:04:54.669278Z",
     "shell.execute_reply": "2022-05-31T03:04:54.668610Z",
     "shell.execute_reply.started": "2022-05-31T03:04:54.663897Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 33 µs, sys: 2 µs, total: 35 µs\n",
      "Wall time: 39.6 µs\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import Dataset\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import os\n",
    "\n",
    "import torchvision.transforms as transforms\n",
    "from torchvision.utils import save_image\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "029bc639-dace-4661-875b-ec0605f5f8e2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-31T03:04:54.670800Z",
     "iopub.status.busy": "2022-05-31T03:04:54.670553Z",
     "iopub.status.idle": "2022-05-31T03:04:54.678709Z",
     "shell.execute_reply": "2022-05-31T03:04:54.678033Z",
     "shell.execute_reply.started": "2022-05-31T03:04:54.670776Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GPU:False\n"
     ]
    }
   ],
   "source": [
    "class Parse:\n",
    "    def __init__(self):  #定义超参数\n",
    "        self.channels = 1\n",
    "        self.img_size = 256\n",
    "        self.latent_dim = 100\n",
    "        self.epochs = 200\n",
    "        self.batch_size = 32\n",
    "        self.lr = 0.001\n",
    "        self.cpu_nums = 8\n",
    "        self.classes = 10\n",
    "        self.want_cuda = \"cuda:0\"\n",
    "        \n",
    "\n",
    "        # 下面是衍生的\n",
    "        self.img_shape = (self.channels, self.img_size, self.img_size)\n",
    "        self.cuda = torch.cuda.is_available()\n",
    "        self.device = self.want_cuda if self.cuda else 'cpu'\n",
    "        self.big_cuda = True  # 能不能把数据集全部放入GPU\n",
    "        print(f\"GPU:{self.cuda}\")\n",
    "opt = Parse()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9a9197d7-5606-4b26-804b-c4e64a5f1354",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-31T03:04:54.779385Z",
     "iopub.status.busy": "2022-05-31T03:04:54.778582Z",
     "iopub.status.idle": "2022-05-31T03:04:54.784634Z",
     "shell.execute_reply": "2022-05-31T03:04:54.784083Z",
     "shell.execute_reply.started": "2022-05-31T03:04:54.779349Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Tools:\n",
    "    # 装一些杂七杂八的工具\n",
    "    def block(in_feat, out_feat, normalize=True):\n",
    "        layers = [nn.Linear(in_feat, out_feat),nn.LeakyReLU(0.2, inplace=True)]\n",
    "        if normalize:\n",
    "            layers.insert(1, nn.BatchNorm1d(out_feat, 0.8))\n",
    "        return layers\n",
    "    \n",
    "    \n",
    "    def sample_image(n_row, batches_done):\n",
    "\n",
    "        z = torch.randn(n_row ** 2, opt.latent_dim)\n",
    "        labels = torch.LongTensor([num for _ in range(n_row) for num in range(n_row)])\n",
    "        gen_imgs = generator(z, labels)\n",
    "        save_image(gen_imgs.data, \"images/%d.png\" % batches_done, nrow=n_row, normalize=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c8d46030-6bfa-40f4-a98e-9208c741d7b9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-31T03:04:54.975866Z",
     "iopub.status.busy": "2022-05-31T03:04:54.975258Z",
     "iopub.status.idle": "2022-05-31T03:04:54.982384Z",
     "shell.execute_reply": "2022-05-31T03:04:54.981283Z",
     "shell.execute_reply.started": "2022-05-31T03:04:54.975840Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Generator(nn.Module):\n",
    "    def __init__(self):\n",
    "        \n",
    "        super().__init__()\n",
    "        self.label_emb = nn.Embedding(opt.classes, opt.classes)\n",
    "        \n",
    "        self.model = nn.Sequential(\n",
    "            *Tools.block(opt.latent_dim + opt.classes, 128, normalize=False),\n",
    "            *Tools.block(128, 256),\n",
    "            *Tools.block(256, 512),\n",
    "            *Tools.block(512, 1024),\n",
    "            nn.Linear(1024, int(np.prod(opt.img_shape))),\n",
    "            nn.Tanh())\n",
    "        \n",
    "        self.optimizer = torch.optim.Adam(self.parameters(), lr=opt.lr)\n",
    "        \n",
    "    def forward(self, noise, labels):\n",
    "        gen_input = torch.cat((self.label_emb(labels), noise), -1)\n",
    "        img = self.model(gen_input)\n",
    "        img = img.view(img.size(0), *opt.img_shape)\n",
    "        return img\n",
    "        \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9724604b-1ac2-4ceb-a049-bb826b26aec3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-31T03:04:55.201203Z",
     "iopub.status.busy": "2022-05-31T03:04:55.200601Z",
     "iopub.status.idle": "2022-05-31T03:04:55.207231Z",
     "shell.execute_reply": "2022-05-31T03:04:55.206688Z",
     "shell.execute_reply.started": "2022-05-31T03:04:55.201178Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Discriminator(nn.Module):\n",
    "    def __init__(self):\n",
    "        \n",
    "        super().__init__()\n",
    "        \n",
    "        self.label_embedding = nn.Embedding(opt.classes, opt.classes)\n",
    "\n",
    "        self.model = nn.Sequential(\n",
    "            nn.Linear(opt.classes + int(np.prod(opt.img_shape)), 512),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            nn.Linear(512, 512),\n",
    "            nn.Dropout(0.4),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            nn.Linear(512, 512),\n",
    "            nn.Dropout(0.4),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            nn.Linear(512, 1),\n",
    "        )\n",
    "        self.optimizer = torch.optim.Adam(self.parameters(), lr=opt.lr)\n",
    "        \n",
    "    def forward(self, img, labels):\n",
    "        \n",
    "        d_in = torch.cat((img.view(img.size(0), -1), self.label_embedding(labels)), -1)\n",
    "        validity = self.model(d_in)\n",
    "        return validity\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d3f43013-2776-4506-8ada-ea7b2674dc64",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-31T03:04:55.663037Z",
     "iopub.status.busy": "2022-05-31T03:04:55.662759Z",
     "iopub.status.idle": "2022-05-31T03:04:55.676440Z",
     "shell.execute_reply": "2022-05-31T03:04:55.675824Z",
     "shell.execute_reply.started": "2022-05-31T03:04:55.663015Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Dev:\n",
    "    def __init__(self):\n",
    "        self.g = Generator()\n",
    "        self.d = Discriminator()\n",
    "        self.loss = nn.MSELoss()\n",
    "        \n",
    "        self.models = (self.g, self.d)\n",
    "        \n",
    "    def load_data(self):\n",
    "        os.makedirs(\"data/mnist\", exist_ok=True)\n",
    "        self.data = torch.utils.data.DataLoader(\n",
    "            datasets.MNIST(\n",
    "                \"../../data/mnist\",\n",
    "                train=True,\n",
    "                download=True,\n",
    "                transform=transforms.Compose(\n",
    "                    [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]\n",
    "                )),\n",
    "            batch_size=opt.batch_size,\n",
    "            shuffle=True,\n",
    "            num_workers=opt.cpu_nums,\n",
    "            drop_last=True)\n",
    "        \n",
    "    def to_cuda(self):\n",
    "        opt.cuda = True\n",
    "        opt.device = opt.want_cuda\n",
    "        for i in self.models:\n",
    "            i.cuda()\n",
    "        if opt.big_cuda:\n",
    "            self.data = torch.utils.data.DataLoader(\n",
    "                self.data.cuda(),\n",
    "                batch_size=opt.batch_size,\n",
    "                shuffle=True,\n",
    "                # num_workers=opt.cpu_nums,\n",
    "                drop_last=True)\n",
    "       \n",
    "    def to_cpu(self):\n",
    "        opt.cuda = False\n",
    "        opt.device = 'cpu'\n",
    "        for i in self.models:\n",
    "            i.cpu()\n",
    "        self.data = torch.utils.data.DataLoader(\n",
    "            self.data,\n",
    "            batch_size=opt.batch_size,\n",
    "            shuffle=True,\n",
    "            num_workers=opt.cpu_nums,\n",
    "            drop_last=True)\n",
    "        \n",
    "        \n",
    "    def train(self):\n",
    "        \n",
    "        \n",
    "        \n",
    "        for i, (imgs, labels) in enumerate(self.data):\n",
    "            \n",
    "            print(i)\n",
    "            \n",
    "            batch_size = imgs.shape[0]\n",
    "            valid = torch.ones(batch_size, 1, device=opt.device)\n",
    "            fake = torch.zeros(batch_size, 1, device=opt.device)\n",
    "            \n",
    "            self.g.optimizer.zero_grad()\n",
    "            z = torch.randn(batch_size, opt.latent_dim)\n",
    "            gen_labels = torch.randint(0, opt.classes, (batch_size,))\n",
    "            \n",
    "            gen_imgs = self.g(z, gen_labels)\n",
    "            validity = self.d(gen_imgs, gen_labels)\n",
    "            g_loss = self.loss(validity, valid)\n",
    "            g_loss.backward()\n",
    "            self.g.optimizer.step()\n",
    "        \n",
    "        \n",
    "        \n",
    "            self.d.optimizer.zero_grad()\n",
    "            validity_real = self.d(imgs, labels)\n",
    "            d_real_loss = self.loss(validity_real, valid)\n",
    "\n",
    "            validity_fake = self.d(gen_imgs.detach(), gen_labels)\n",
    "            d_fake_loss = self.loss(validity_fake, fake)\n",
    "\n",
    "            d_loss = (d_real_loss + d_fake_loss) / 2\n",
    "            d_loss.backward()\n",
    "            self.d.optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dab4da40-099f-4b8e-a2e5-735901a9ceb2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
